import os
import torch
import argparse
from model import ActionClassificationModel
from dataset import TestDataset
from torchvision.transforms import transforms

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def predict(model, image_folder):
    transform = transforms.Compose([
        transforms.Resize((112, 112)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    dataset = TestDataset(image_folder, transform=transform)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)

    model.eval()
    classes = dataset.classes

    with torch.no_grad():
        for i, (inputs, _) in enumerate(dataloader):
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            predicted_label = classes[predicted.item()]
            print("Image {}: {}".format(i, predicted_label))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Action Classification Predict')
    parser.add_argument('--model-path', type=str, default="pth/model_epoch_300.pth",
                        help='path to the trained model')
    parser.add_argument('--image-folder', type=str, default="dataset/test/3",
                        help='path to the folder containing images')
    args = parser.parse_args()

    model_path = args.model_path
    image_folder = args.image_folder

    if not os.path.exists(model_path):
        print("Model file not found at {}".format(model_path))
    elif not os.path.isdir(image_folder):
        print("Image folder not found at {}".format(image_folder))
    else:
        model = ActionClassificationModel()
        model.load_state_dict(torch.load(model_path,map_location=device))
        predict(model, image_folder)
